#  -*- coding: utf-8 -*-

import traceback
import matplotlib.pyplot as plt
import numpy as np
from pymongo import UpdateOne
import pandas as pd
from data.data_module import DataModule
from trading.signal.computer.base_signal_computer import BaseSignalComputer
from util.stock_util import get_all_codes
import mpl_finance as mpf


class BollSignalComputer(BaseSignalComputer):
    def __init__(self):
        BaseSignalComputer.__init__(self, 'boll_signal')

    def compute(self, begin_date, end_date):
        """
        计算指定日期内的信号
        :param begin_date: 开始日期
        :param end_date: 结束日期
        """
        # all_codes = get_all_codes('2018-08-24')
        all_codes = ['601398']
        dm = DataModule()

        N = 20
        k = 2
        mistake = 0.1
        for code in all_codes[:]:
            print(code)
            try:
                df_daily = dm.get_k_data(code, begin_date=begin_date, autype='kba', end_date=end_date, trade=True)
                df_daily['value'] = df_daily['close'].astype('float') - df_daily['close'].astype('float').shift(1)
                df_daily['value1'] = df_daily['value']
                df_daily['value1'][df_daily['value1'] < 0] = 0
                df_daily['value2'] = df_daily['value']
                df_daily['value2'][df_daily['value1'] > 0] = 0
                # print (df_daily['value1'])
                # print(df_daily['value2'])
                # df_daily['plus']=df_daily['value1'].rolling(24).max()
                # df_daily['minus'] = df_daily['value2'].rolling(24).min()

                # print (df_daily['value1'].rolling(9).sum())
                df_daily['plus6'] = df_daily['value1'].rolling(6).sum()
                df_daily['minus6'] = abs(df_daily['value2'].rolling(6).sum())


                # print ( df_daily['plus6'].tolist())
                # print(df_daily['minus6'].tolist())
                df_daily['rsi6'] = df_daily['plus6'] / (df_daily['plus6'] + df_daily['minus6'])*100

                # print(df_daily['rsi'])
                # print (df_daily['plus'])
                # # print ( df_daily['minus'])
                # df_daily['rsi']=df_daily['plus']/(df_daily['plus']-df_daily['minus'])*100
                # print (df_daily['rsi'])
                print (df_daily)
                plt.plot(df_daily['rsi6'])
                # plt.plot(df_daily['rsi12'])
                # plt.plot(df_daily['rsi24'])
                # plt.plot(df_daily['d'])
                # plt.plot(df_daily['j'])
                plt.xticks([])
                plt.show()
            except:
                traceback.print_exc()
                # break


if __name__ == '__main__':
    BollSignalComputer().compute(begin_date='2018-01-01', end_date='2018-08-30')
